import torch
import numpy as np


def bbox_iou(box1, box2):
    '''
    The size of box1 is bs*wid*hei*2*4
    The size of box2 is bs*2*4
    '''
    
    inter_rect_x1 = torch.max(box1[0,0], box2[0,0])
    inter_rect_y1 = torch.min(box1[1,0], box2[1,0])
    inter_rect_x2 = torch.min(box1[0,1], box2[0,1])
    inter_rect_y2 = torch.min(box1[1,1], box2[1,1])
    inter_rect_x3 = torch.max(box1[0,2], box2[0,2])
    inter_rect_y3 = torch.max(box1[1,2], box2[1,2])
    inter_rect_x4 = torch.min(box1[0,3], box2[0,3])
    inter_rect_y4 = torch.max(box1[1,3], box2[1,3])
    rect_1 = np.array([[1,1,1],[inter_rect_x1, inter_rect_x2, inter_rect_x3],[inter_rect_y1, inter_rect_y2, inter_rect_y3]])
    rect_1_area = np.linalg.det(rect_1)/2
    rect_2 = np.array([[1,1,1],[inter_rect_x2, inter_rect_x3, inter_rect_x4],[inter_rect_y2, inter_rect_y3, inter_rect_y4]])
    rect_2_area = np.linalg.det(rect_2)/2
    inter_area = rect_1_area + rect_2_area
    box1_area = (box1[0,1] - box1[0,0]) * (box1[1,0] - box1[1,3])
    box2_area = (box2[0,1] - box2[0,0]) * (box2[1,0] - box2[1,3])
    iou = inter_area/(box1_area + box2_area - inter_area)
    return iou


def intector_mn(T_mn, A_mn):
    '''
    Select the best location
    '''
    Thre = 0.3
    mask = torch.zeros(T_mn.size(0), T_mn.size(1), T_mn.size(2))
    for i in range(T_mn.size(0)):
        for m in range(T_mn.size(1)):
            for n in range(T_mn.size(2)):
                if bbox_iou(T_mn[i,m,n,:,:], A_mn[i,:,:]) > Thre:
                    mask[i,m,n] = 1
    return mask


def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)


def non_max_suppression(prediction):
    '''
    找到v1值最大的位置，每张图片中只有一个车牌
    :param prediction:bs*M*N*8
    :return:
    '''
    bs = prediction.shape[0]
    print('Num of bs is %d' %bs)
    v1 = prediction[:,:,:,0]
    out = [None for _ in range(bs)]
    for i in range(bs):
        cur_v1 = v1[i, :, :]
        max_row, index_col = torch.max(cur_v1, 1)
        index_row = torch.max(max_row,0)[1].item()
        index_col = index_col[index_row].item()
        out[i] = prediction[i, index_row, index_col]
    print('out is {}'.format(out))
    return out

